package com.sxd.bill.ner;

import java.io.File;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Date;
import java.util.List;
import java.util.Map;

import com.hankcs.hanlp.corpus.io.IOUtil;
import com.hankcs.hanlp.model.crf.CRFNERecognizer;
import com.hankcs.hanlp.model.crf.crfpp.Encoder;
import com.hankcs.hanlp.model.crf.crfpp.crf_learn;
import com.hankcs.hanlp.model.perceptron.utility.Utility;
import com.sxd.bill.common.item.TextLabelItem;
import com.sxd.bill.ner.utils.OutputUtility;
import com.sxd.bill.ner.utils.WordUtitly;

public class CRFNerWrapper {
    private CRFNERecognizer crf;

    public CRFNerWrapper(String modelPath) throws Exception {
        this.crf = new CRFNERecognizer(modelPath);
    }

    public String[] recognize(String[] words, String[] kinds) {
        String[] preds = crf.recognize(words, kinds);
        return preds;
    }

    public String[] recognize(String text, String[] kinds) {
        String[] preds = recognize(text.split(""), kinds);
        return preds;
    }

    public String[] recognize(String text) {
        // // 因为没使用分词工具，故每个word没有词性标签，此处词性标签统一设为"unk"(与训练数据一致)
        // String[] kinds = new String[text.length()];
        // for (int i = 0; i < kinds.length; i++) {
        // kinds[i] = "unk";
        // }
        // 获取每个字符的类别（中文、英文、数字、特征符号、其他）
        String[] kinds = WordUtitly.getCharKind(text);
        // 识别实体标签
        String[] preds = recognize(text, kinds);
        return preds;
    }

    public TextLabelItem recognizeLabelEntities(String text) {
        String[] tags = recognize(text);
        TextLabelItem textLabelItem = OutputUtility.getTextLabelItem(Arrays.asList(tags), text);
        return textLabelItem;
    }

    public List<TextLabelItem> recognizeLabelEntities(List<String> texts) {
        List<TextLabelItem> textLabelItems = new ArrayList<>();
        for (String text : texts) {
            TextLabelItem textLabelItem = recognizeLabelEntities(text);
            textLabelItems.add(textLabelItem);
        }
        return textLabelItems;
    }

    /**
     * 特征模板
     * 
     * @return
     */
    public static String getDefaultFeatureTemplate() {
        return "# Unigram\nU0:%x[-2,0]\nU1:%x[-1,0]\nU2:%x[0,0]\nU3:%x[1,0]\nU4:%x[2,0]\nU5:%x[-2,1]\nU6:%x[-1,1]\nU7:%x[0,1]\nU8:%x[1,1]\nU9:%x[2,1]\nUA:%x[-2,1]%x[-1,1]\nUB:%x[-1,1]%x[0,1]\nUC:%x[0,1]%x[1,1]\nUD:%x[1,1]%x[2,1]\nUE:%x[2,1]%x[3,1]\n\n# Bigram\nB";
    }

    /**
     * 模型训练
     * 
     * @param trainPath
     * @param templatePath
     * @param modelPath
     * @param maxIter
     * @throws Exception
     */
    public static void train(String trainPath, String templatePath, String modelPath, int maxIter) throws Exception {
        CRFNERecognizer crf = new CRFNERecognizer(null);
        crf_learn.Option option = new crf_learn.Option();
        crf.train(templatePath, trainPath, modelPath, maxIter, option.freq, option.eta, option.cost,
                option.thread, option.shrinking_size, Encoder.Algorithm.fromString(option.algorithm));
    }

    public static void train(String trainPath, String modelPath, int maxIter) throws Exception {
        // 使用默认的特征模板
        String templFile = null;
        File tmpTemplate = File.createTempFile("crfpp-template-" + new Date().getTime(), ".txt");
        tmpTemplate.deleteOnExit();
        templFile = tmpTemplate.getAbsolutePath();
        String template = getDefaultFeatureTemplate();
        IOUtil.saveTxt(templFile, template);

        train(trainPath, templFile, modelPath, maxIter);
    }

    /**
     * 模型效果评估(TODO)
     * 
     * @param modelPath
     * @param evalPath
     */
    public static void eval(String modelPath, String evalPath) throws Exception {
        CRFNERecognizer crf = new CRFNERecognizer(modelPath);
        Map<String, double[]> score = Utility.evaluateNER(crf, evalPath);
        Utility.printNERScore(score);
    }

}
